{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part 1, Topic 2: 1.5 Round AES Attack"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "NOTE: This lab references some (commercial) training material on [ChipWhisperer.io](https://www.ChipWhisperer.io). You can freely execute and use the lab per the open-source license (including using it in your own courses if you distribute similarly), but you must maintain notice about this source location. Consider joining our training course to enjoy the full experience.\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**SUMMARY:** *In the last lab, we made a small modification that allowed us to mount our simple AES round skip fault. In the unmodified code, however, the lack of a check at the beginning of the loop prevented us from being able to skip the first round in AES. In this lab, we'll break the original code using an attack generated by https://github.com/cbouilla/AES-attacks-finder.*\n",
    "\n",
    "**LEARNING OUTCOMES:**\n",
    "* Identifying the end of the first AES round"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SCOPETYPE = 'OPENADC'\n",
    "PLATFORM = 'CWLITEARM'\n",
    "CRYPTO_TARGET='TINYAES128C'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash -s \"$PLATFORM\" \"$CRYPTO_TARGET\"\n",
    "cd ../../../firmware/mcu/simpleserial-aes\n",
    "make PLATFORM=$1 CRYPTO_TARGET=$2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run \"../Setup_Scripts/Setup_Generic.ipynb\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fw_path = \"../../../firmware/mcu/simpleserial-aes/simpleserial-aes-{}.hex\".format(PLATFORM)\n",
    "cw.program_target(scope, prog, fw_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if PLATFORM == \"CWLITEXMEGA\":\n",
    "    def reboot_flush():            \n",
    "        scope.io.pdic = False\n",
    "        time.sleep(0.1)\n",
    "        scope.io.pdic = \"high_z\"\n",
    "        time.sleep(0.1)\n",
    "        #Flush garbage too\n",
    "        target.flush()\n",
    "else:\n",
    "    def reboot_flush():            \n",
    "        scope.io.nrst = False\n",
    "        time.sleep(0.05)\n",
    "        scope.io.nrst = \"high_z\"\n",
    "        time.sleep(0.05)\n",
    "        #Flush garbage too\n",
    "        target.flush()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Again, we'll start by collecting a reference ciphertext and output:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scope.clock.adc_src = \"clkgen_x1\"\n",
    "reboot_flush()\n",
    "scope.arm()\n",
    "target.simpleserial_write('p', bytearray([0]*16))\n",
    "ret = scope.capture()\n",
    "if ret:\n",
    "    print(\"No trigger!\")\n",
    "\n",
    "wave = scope.get_last_trace()\n",
    "\n",
    "output = target.simpleserial_read_witherrors('r', 16)\n",
    "gold_ct = output['payload']\n",
    "\n",
    "print(gold_ct)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib notebook\n",
    "import matplotlib.pyplot as plt\n",
    "plt.figure()\n",
    "plt.plot(wave[:2000])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Identify a good glitch range. We're still attacking after the AddRoundKey operation, so look at the second of those."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "glitch_loc = range(870, 920)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if scope._is_husky:\n",
    "    scope.glitch.enabled = True\n",
    "scope.glitch.clk_src = \"clkgen\"\n",
    "scope.glitch.output = \"clock_xor\"\n",
    "scope.glitch.trigger_src = \"ext_single\"\n",
    "scope.glitch.repeat = 1\n",
    "scope.io.hs2 = \"glitch\"\n",
    "# These width/offset numbers are for CW-Lite/Pro; for CW-Husky, convert as per Fault 1_1:\n",
    "scope.glitch.width = 3\n",
    "scope.glitch.offset = -12.8\n",
    "print(scope.glitch)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Again, we'll be using power analysis to determine when we've got a good glitch. This time you should see the first round being completed:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm.notebook import tqdm, trange\n",
    "wave = None\n",
    "import logging\n",
    "ktp = cw.ktp.Basic()\n",
    "logging.getLogger().setLevel(logging.ERROR)\n",
    "reboot_flush()\n",
    "for i in trange(min(glitch_loc), max(glitch_loc) + 1):\n",
    "    scope.adc.timeout = 0.2\n",
    "    scope.glitch.ext_offset = i\n",
    "    ack = None\n",
    "    while ack is None:\n",
    "        target.simpleserial_write('k', ktp.next()[0])\n",
    "        ack = target.simpleserial_wait_ack()\n",
    "        if ack is None:\n",
    "            reboot_flush()\n",
    "            time.sleep(0.1)\n",
    "    \n",
    "    scope.arm()\n",
    "    \n",
    "    pt = bytearray([0]*16)\n",
    "    target.simpleserial_write('p', pt)\n",
    "    ret = scope.capture()\n",
    "    if ret:\n",
    "        reboot_flush() #bad if we accidentally didn't have this work\n",
    "        time.sleep(0.1)\n",
    "        print(\"timed out!\")\n",
    "        continue\n",
    "    output = target.simpleserial_read_witherrors('r', 16, glitch_timeout = 1)\n",
    "    if output['valid']:\n",
    "        if output['payload'] != gold_ct:\n",
    "            print(\"Glitched at {}\".format(i))\n",
    "            wave = scope.get_last_trace()\n",
    "            break\n",
    "    else:\n",
    "        reboot_flush()\n",
    "        \n",
    "%matplotlib notebook\n",
    "import matplotlib.pyplot as plt\n",
    "plt.figure()\n",
    "plt.plot(wave)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "glitched_ct0 = bytearray(output['payload'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "glitched_ct0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As it turns out, there's actually two possible `glitched_ct0`s you could've gotten here. One is adding the second round key at the end of the encryption and the other is adding the 0th round key! This actually makes the attack a bit simpler (and much faster), as we can recover the key with one plaintext and both of the possible faulty ciphertexts.\n",
    "\n",
    "Let's rerun our loop until we get another glitch that breaks out early, but is different than our first glitch:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm.notebook import tqdm, trange\n",
    "wave = None\n",
    "import logging\n",
    "ktp = cw.ktp.Basic()\n",
    "logging.getLogger().setLevel(logging.ERROR)\n",
    "reboot_flush()\n",
    "while True:\n",
    "    scope.adc.timeout = 0.2\n",
    "    scope.glitch.ext_offset = i\n",
    "    ack = None\n",
    "    while ack is None:\n",
    "        target.simpleserial_write('k', ktp.next()[0])\n",
    "        ack = target.simpleserial_wait_ack()\n",
    "        if ack is None:\n",
    "            reboot_flush()\n",
    "            time.sleep(0.1)\n",
    "    \n",
    "    scope.arm()\n",
    "    \n",
    "    pt = bytearray([0]*16)\n",
    "    target.simpleserial_write('p', pt)\n",
    "    ret = scope.capture()\n",
    "    if ret:\n",
    "        reboot_flush() #bad if we accidentally didn't have this work\n",
    "        time.sleep(0.1)\n",
    "        print(\"timed out!\")\n",
    "        continue\n",
    "    output = target.simpleserial_read_witherrors('r', 16, glitch_timeout = 1)\n",
    "    if output['valid']:\n",
    "        if output['payload'] != gold_ct:\n",
    "            if output['payload'] != glitched_ct0:\n",
    "                print(\"Glitched at {}\".format(i))\n",
    "                wave = scope.get_last_trace()\n",
    "                break\n",
    "    else:\n",
    "        reboot_flush()\n",
    "        \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib notebook\n",
    "import matplotlib.pyplot as plt\n",
    "plt.figure()\n",
    "plt.plot(wave)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(output['payload'])\n",
    "glitched_ct1 = bytearray(output['payload'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now feed it into our autogenerated solver. `Known` is just a big array of known information (includes both ciphertexts and the plaintext, which is just all 0's in our case):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Known = [0] * 49\n",
    "Known[0] = 1\n",
    "\n",
    "for i in range(4):\n",
    "    for j in range(4):\n",
    "        Known[48-(i*4+j)] = glitched_ct0[i+j*4]\n",
    "        Known[16-(i*4+j)] = glitched_ct1[i+j*4]\n",
    "        \n",
    "from out2 import MakeTableMul2_8, Attack\n",
    "\n",
    "MakeTableMul2_8()\n",
    "kguess = Attack(Known)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you don't get a solution, you'll probably just need to swap `glitched_ct0` and `glitched_ct1` (the order does matter, but we don't know which is which). We can print our solution in a nicer format with the following:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(4):\n",
    "    print([hex(kguess[i][j]) for j in range(4)])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Variations on this Attack\n",
    "\n",
    "This attack is pretty specific to this implementation. What if, for example, we were able to break out at the same spot, but only the 2nd round key was used for the final AddRoundKey? In that case, we could run an attack with two plaintext faulted ciphertexts pairs, though it is much slower than the attack we did. Different situations will have different requirements, though attacks quickly become infeasible the further you get into AES (as you'd expect). The tool we used to generate the attack for this situation is available at https://github.com/cbouilla/AES-attacks-finder. If you're curious, try thinking up different glitch scenarios and use the tools to see if the attacks are possible! The tool outputs C code. In the case of this attack, the generated C code was ported to python with the help of `ctopy`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
